// MobileCapture.cpp : Defines the entry point for the application.

#include "stdafx.h"
#include "MobileCapture.h"

#include <regext.h>

#define HOTKEY_ID 400

typedef struct __HardwareKey {
	UINT modifiers;
	UINT vk;
} HardwareKey;

static HardwareKey gbl_hardwareKeys[64];

// Save the bitmap 'hbmpSource' into the file specified by 'pszFilename'
static void SaveBitmap(WCHAR * pszFilename, HDC hdcSource, HBITMAP hbmpSource)
{
	// Determine the size of the source bitmap
	BITMAP bm;
	GetObject(hbmpSource, sizeof(BITMAP), &bm);

	// Create a device independant bitmap (DIB)
	// that is the same size
	BITMAPINFO bmi;
	memset(&bmi, 0, sizeof(bmi));
	bmi.bmiHeader.biHeight = bm.bmHeight;
	bmi.bmiHeader.biWidth = bm.bmWidth; 
	bmi.bmiHeader.biPlanes = 1;
	bmi.bmiHeader.biBitCount = 24;
	bmi.bmiHeader.biSize = sizeof(bmi.bmiHeader);
	int ls = (bmi.bmiHeader.biWidth + 3) & (~3);
	bmi.bmiHeader.biSizeImage = ls * bmi.bmiHeader.biHeight * bmi.bmiHeader.biBitCount / 8;
	bmi.bmiHeader.biXPelsPerMeter = 3780;
	bmi.bmiHeader.biYPelsPerMeter = 3780;

	void *pBuffer;
	HBITMAP hTargetBitmap = CreateDIBSection(hdcSource, &bmi, DIB_RGB_COLORS, (void**)&pBuffer, NULL, 0);

	// Copy the source bitmap into the DIB.
	HDC hdcDest = CreateCompatibleDC(hdcSource);
	HBITMAP hOldBitmap = (HBITMAP)SelectObject(hdcDest, hTargetBitmap);
	BitBlt(hdcDest, 0, 0, bm.bmWidth, bm.bmHeight, hdcSource, 0, 0, SRCCOPY);
	SelectObject(hdcDest, hOldBitmap);
	DeleteDC(hdcDest);

	// Finally save the bits to disk
	HANDLE hFile = CreateFile(pszFilename, GENERIC_WRITE, 0, NULL, CREATE_ALWAYS, FILE_ATTRIBUTE_NORMAL, NULL);
	if (hFile != INVALID_HANDLE_VALUE)
	{
		BITMAPFILEHEADER hdr;
		hdr.bfType		= ((WORD) ('M' << 8) | 'B');
		hdr.bfSize		= sizeof(hdr) + sizeof(BITMAPINFOHEADER);
		hdr.bfReserved1 = 0;
		hdr.bfReserved2 = 0;
		hdr.bfOffBits	= hdr.bfSize;
		
		// Write the BITMAPFILEHEADER and BITMAPINFOHEADER
		// structures followed by the raw pixel data
		DWORD dwByteCount;
		WriteFile(hFile, &hdr, sizeof(hdr), &dwByteCount, NULL);
		WriteFile(hFile, &bmi.bmiHeader, sizeof(bmi.bmiHeader), &dwByteCount, NULL);
		WriteFile(hFile, pBuffer, bmi.bmiHeader.biSizeImage, &dwByteCount, NULL);

		CloseHandle(hFile);
	}

	// Free the bitmap resource
	DeleteObject(hTargetBitmap);
}


// Take a snapshot of a region of the screen and save it to disk
static void Snapshot(WCHAR * pszFilename, RECT *pRect)
{
	// Create a new bitmap of the required size and a device
	// context to allow us to draw to it.
	int width = pRect->right - pRect->left;
	int height = pRect->bottom - pRect->top;

	HDC hdcSrc = GetDC(NULL);
	HDC hdcDest = CreateCompatibleDC(hdcSrc);
	HBITMAP hbmp = CreateCompatibleBitmap(hdcSrc, width, height);

	SelectObject(hdcDest, hbmp);

	// Blit (copy) from the source device context (the screen)
	// to the device context that is associated with our
	// offscreen bitmap buffer.
	BitBlt(hdcDest, 0, 0, width, height, hdcSrc, pRect->left, pRect->top, SRCCOPY);

	// Finally save our bitmap to disk
	SaveBitmap(pszFilename, hdcDest, hbmp);

	// Free up the resources
	DeleteDC(hdcDest);
	DeleteObject(hbmp);

	ReleaseDC(NULL, hdcSrc);
}

static void CaptureImage(HWND hDlg, UINT i)
{
	int x1 = GetDlgItemInt(hDlg, IDC_EDIT_X1, NULL, FALSE);
	int y1 = GetDlgItemInt(hDlg, IDC_EDIT_Y1, NULL, FALSE);
	int x2 = GetDlgItemInt(hDlg, IDC_EDIT_X2, NULL, FALSE);
	int y2 = GetDlgItemInt(hDlg, IDC_EDIT_Y2, NULL, FALSE);

	WCHAR szFileName[MAX_PATH];
	swprintf(szFileName, L"\\screenshot%03d.bmp", i);

	RECT rcBounds = {x1, y1, x2, y2};
	Snapshot(szFileName, &rcBounds);

	PlaySound(MAKEINTRESOURCE(IDR_WAVE_CAPTURE), GetModuleHandle(NULL), SND_ASYNC | SND_RESOURCE | SND_NODEFAULT);
}

// Look within the HKLM\Software\Microsoft\Shell\Keys registry key
// for a list of hardware keys that can be used as triggers to start
// the screen capture process.
static void PopulateHardwareKeys(HWND hDlg)
{
	HKEY hKey;

	if (RegOpenKeyEx(HKEY_LOCAL_MACHINE, L"Software\\Microsoft\\Shell\\Keys", 0, 0, &hKey) == ERROR_SUCCESS)
	{
		int i = 0, index = 0;
		WCHAR szKey[MAX_PATH];
		DWORD dwKeyLen = MAX_PATH;

		// Loop over each key
		while (RegEnumKeyEx(hKey, i++, szKey, &dwKeyLen, NULL, NULL, NULL, NULL) == ERROR_SUCCESS)
		{
			WCHAR szName[MAX_PATH];
			if (RegistryGetString(hKey, szKey, L"Name", szName, MAX_PATH) == ERROR_SUCCESS)
			{
				// Add the key's name to the combo box
				SendDlgItemMessage(hDlg, IDC_COMBO_HARDWAREKEY, CB_ADDSTRING, 0, (LPARAM)szName);

				// Convert the registry key name (i.e. 40C1)
				// into the two values we need to pass to
				// RegisterHotKey (MOD_WIN and VK_APP1) and
				// then store them in the gbl_hardwareKeys
				// array.
				gbl_hardwareKeys[index].vk = wcstol(szKey, NULL, 16);
				gbl_hardwareKeys[index].modifiers = 0;

				if (gbl_hardwareKeys[index].vk & 0x4000)
					gbl_hardwareKeys[index].modifiers |= MOD_WIN;
				if (gbl_hardwareKeys[index].vk & 0x2000)
					gbl_hardwareKeys[index].modifiers |= MOD_ALT;
				if (gbl_hardwareKeys[index].vk & 0x0800)
					gbl_hardwareKeys[index].modifiers |= MOD_CONTROL;
				if (gbl_hardwareKeys[index].vk & 0x0400)
					gbl_hardwareKeys[index].modifiers |= MOD_SHIFT;

				gbl_hardwareKeys[index++].vk &= 0xFF;
			}

			dwKeyLen = MAX_PATH;
		}

		RegCloseKey(hKey);
	}

	// Select the first hardware key by default
	SendDlgItemMessage(hDlg, IDC_COMBO_HARDWAREKEY, CB_SETCURSEL, 0, 0); 
}

INT_PTR CALLBACK MobileCaptureDialog(HWND hDlg, UINT message, WPARAM wParam, LPARAM lParam)
{
	HWND hWndCtrl;
	BYTE appkey;

    switch (message)
    {
        case WM_INITDIALOG:
            {
                // Create a Done button and size it.  
                SHINITDLGINFO shidi;
				SHMENUBARINFO mbi;

                shidi.dwMask = SHIDIM_FLAGS;
                shidi.dwFlags = SHIDIF_DONEBUTTON | SHIDIF_SIPDOWN | SHIDIF_SIZEDLGFULLSCREEN | SHIDIF_EMPTYMENU;
                shidi.hDlg = hDlg;
                SHInitDialog(&shidi);

				// Create a menu bar
				memset(&mbi, 0, sizeof(SHMENUBARINFO));
				mbi.cbSize     = sizeof(SHMENUBARINFO);
				mbi.hwndParent = hDlg;
				mbi.nToolBarId = IDR_MENU;
				mbi.hInstRes   = GetModuleHandle(NULL);
				mbi.dwFlags = SHCMBF_HMENU;
				SHCreateMenuBar(&mbi);

				// Set up the default selections, making sure our selection
				// changed processing is triggered to update the UI controls.
				SendDlgItemMessage(hDlg, IDC_RADIO_FULLSCREEN, BM_SETCHECK, BST_CHECKED, 0);
				SendMessage(hDlg, WM_COMMAND, IDC_RADIO_FULLSCREEN, 0);
				SendDlgItemMessage(hDlg, IDC_RADIO_HARDWAREKEY, BM_SETCHECK, BST_CHECKED, 0);
				SendMessage(hDlg, WM_COMMAND, IDC_RADIO_HARDWAREKEY, 0);

				// Set the range and current value of the spinboxes
				SendDlgItemMessage(hDlg, IDC_SPIN_DELAYTIMER, UDM_SETRANGE32, 1, 60);
				SendDlgItemMessage(hDlg, IDC_SPIN_DELAYTIMER, UDM_SETPOS, 0, MAKELPARAM(1, 0));
				SendDlgItemMessage(hDlg, IDC_SPIN_MULTIPLE, UDM_SETRANGE32, 1, 25);
				SendDlgItemMessage(hDlg, IDC_SPIN_MULTIPLE, UDM_SETPOS, 0, MAKELPARAM(1, 0));

				// Populate the list of hardware keys
				PopulateHardwareKeys(hDlg);
            }
            return (INT_PTR)TRUE;

		case WM_HOTKEY:
			if (wParam == HOTKEY_ID)
			{
				// Capture an image and then shutdown the application
				// in response to the user pressing the hardware key.
				CaptureImage(hDlg, 1);
				EndDialog(hDlg, IDOK);
			}
			break;

        case WM_COMMAND:
            switch (LOWORD(wParam))
			{
				case IDOK:
					EndDialog(hDlg, LOWORD(wParam));
					return (INT_PTR)TRUE;

				case IDC_RADIO_FULLSCREEN:
				case IDC_RADIO_AREA:
					// Enable and disable controls to suit the current
					// state of the fullscreen and area radio buttons.
					if (LOWORD(wParam) == IDC_RADIO_FULLSCREEN)
					{
						SetDlgItemInt(hDlg, IDC_EDIT_X1, 0, FALSE);
						SetDlgItemInt(hDlg, IDC_EDIT_Y1, 0, FALSE);
						SetDlgItemInt(hDlg, IDC_EDIT_X2, GetSystemMetrics(SM_CXSCREEN), FALSE);
						SetDlgItemInt(hDlg, IDC_EDIT_Y2, GetSystemMetrics(SM_CYSCREEN), FALSE);
					}
					hWndCtrl = GetDlgItem(hDlg, IDC_EDIT_X1);
					EnableWindow(hWndCtrl, LOWORD(wParam) == IDC_RADIO_AREA);
					hWndCtrl = GetDlgItem(hDlg, IDC_EDIT_Y1);
					EnableWindow(hWndCtrl, LOWORD(wParam) == IDC_RADIO_AREA);
					hWndCtrl = GetDlgItem(hDlg, IDC_EDIT_X2);
					EnableWindow(hWndCtrl, LOWORD(wParam) == IDC_RADIO_AREA);
					hWndCtrl = GetDlgItem(hDlg, IDC_EDIT_Y2);
					EnableWindow(hWndCtrl, LOWORD(wParam) == IDC_RADIO_AREA);
					break;

				case IDC_RADIO_HARDWAREKEY:
				case IDC_RADIO_DELAYTIMER:
					// Enable and disable controls to suit the current
					// state of the hardware key and delay timer radio
					// buttons.
					if (LOWORD(wParam) == IDC_RADIO_HARDWAREKEY)
					{
						SendDlgItemMessage(hDlg, IDC_CHECK_MULTIPLE, BM_SETCHECK, BST_UNCHECKED, 0);
						SetDlgItemText(hDlg, IDC_EDIT_MULTIPLE, L"1");
					}
					hWndCtrl = GetDlgItem(hDlg, IDC_COMBO_HARDWAREKEY);
					EnableWindow(hWndCtrl, LOWORD(wParam) == IDC_RADIO_HARDWAREKEY);
					hWndCtrl = GetDlgItem(hDlg, IDC_EDIT_DELAYTIMER);
					EnableWindow(hWndCtrl, LOWORD(wParam) == IDC_RADIO_DELAYTIMER);
					hWndCtrl = GetDlgItem(hDlg, IDC_CHECK_MULTIPLE);
					EnableWindow(hWndCtrl, LOWORD(wParam) == IDC_RADIO_DELAYTIMER);
					hWndCtrl = GetDlgItem(hDlg, IDC_EDIT_MULTIPLE);
					EnableWindow(hWndCtrl, LOWORD(wParam) == IDC_RADIO_DELAYTIMER);
					break;

				case IDC_CHECK_MULTIPLE:
					// Enable and disable controls to suit the current
					// state of the multiple capture checkbox.
					hWndCtrl = GetDlgItem(hDlg, IDC_EDIT_MULTIPLE);
					if (SendDlgItemMessage(hDlg, IDC_CHECK_MULTIPLE, BM_GETCHECK, 0, 0) == BST_CHECKED)
					{
						SetWindowText(hWndCtrl, L"2");
						EnableWindow(hWndCtrl, TRUE);
					}
					else
					{
						SetWindowText(hWndCtrl, L"1");
						EnableWindow(hWndCtrl, FALSE);
					}
					break;

				case IDC_CAPTURE:			
					// Hide the main window
					ShowWindow(hDlg, SW_MINIMIZE);

					if (SendDlgItemMessage(hDlg, IDC_RADIO_HARDWAREKEY, BM_GETCHECK, 0, 0) == BST_CHECKED)
					{
						// We're in hardware key capture mode. So register the hotkey which
						// will notify us when the user wants us to capture a screenshot.
						int i = SendDlgItemMessage(hDlg, IDC_COMBO_HARDWAREKEY, CB_GETCURSEL, 0, 0);
						HardwareKey key = gbl_hardwareKeys[i];

						RegisterHotKey(hDlg, HOTKEY_ID, key.modifiers, key.vk);
					}
					else
					{
						// We're in timer capture mode. So determine the number of seconds to
						// delay between captures
						int delay = GetDlgItemInt(hDlg, IDC_EDIT_DELAYTIMER, NULL, FALSE) * 1000;

						// Then capture the desired number of images.
						for (UINT i = 0; i < GetDlgItemInt(hDlg, IDC_EDIT_MULTIPLE, NULL, FALSE); i++)
						{
							Sleep(delay);

							CaptureImage(hDlg, i + 1);
						}

						// And finally shutdown the application
						EndDialog(hDlg, IDOK);
					}
					break;
			}
			break;

		case WM_CLOSE:
            EndDialog(hDlg, message);
            return (INT_PTR)TRUE;

		case WM_DESTROY:
			UnregisterHotKey(hDlg, HOTKEY_ID);
			break;

    }
    return (INT_PTR)FALSE;
}

int WINAPI WinMain(HINSTANCE hInstance,
                   HINSTANCE hPrevInstance,
                   LPTSTR    lpCmdLine,
                   int       nCmdShow)
{
	return DialogBox(hInstance, (LPCTSTR)IDD_MOBILECAPTURE, NULL, MobileCaptureDialog);
}